import os
from typing import List, Dict, Any
import autogen
from chromadb import PersistentClient
from chromadb.utils import embedding_functions
import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
import asyncio

class JobAgentSystem:
    def __init__(self):
        # 初始化配置
        self.config = self.load_config()
        
        # 初始化向量数据库
        self.chroma_client = PersistentClient(path="b/data/chroma_db")
        self.embedding_func = embedding_functions.DefaultEmbeddingFunction()
        
        # 初始化智能体
        self.agents = self.setup_agents()
        
        # 初始化技能向量集合
        self.setup_skill_embeddings()
        
    def load_config(self) -> Dict[str, Any]:
        """加载系统配置"""
        return {
            "llm_config": {
                "config_list": [
                    {
                        "model": "gpt-4",
                        "api_key": os.getenv("OPENAI_API_KEY")
                    }
                ],
                "temperature": 0.7
            },
            "max_agents": 5,
            "skill_threshold": 0.85
        }
    
    def setup_skill_embeddings(self):
        """初始化技能向量数据库"""
        self.skill_collection = self.chroma_client.get_or_create_collection(
            name="skills",
            embedding_function=self.embedding_func
        )
        
        # 添加基础技能
        base_skills = ["Python", "Java", "SQL", "Machine Learning", 
                      "Data Analysis", "Cloud Computing", "Statistics"]
        embeddings = self.embedding_func(base_skills)
        
        self.skill_collection.add(
            embeddings=embeddings,
            documents=base_skills,
            ids=[f"skill_{i}" for i in range(len(base_skills))]
        )
    
    def setup_agents(self) -> Dict[str, autogen.AssistantAgent]:
        """设置多智能体系统"""
        # 分析智能体
        analyst_agent = autogen.AssistantAgent(
            name="Analyst",
            system_message="你是一个数据分析专家，擅长分析岗位需求和市场趋势。",
            llm_config=self.config["llm_config"]
        )
        
        # 匹配智能体
        matcher_agent = autogen.AssistantAgent(
            name="Matcher",
            system_message="你负责匹配求职者技能与岗位需求，计算匹配度。",
            llm_config=self.config["llm_config"]
        )
        
        # 趋势智能体
        trend_agent = autogen.AssistantAgent(
            name="TrendAnalyst",
            system_message="你分析行业趋势和新兴技术需求。",
            llm_config=self.config["llm_config"]
        )
        
        return {
            "analyst": analyst_agent,
            "matcher": matcher_agent,
            "trend": trend_agent
        }
    
    async def analyze_industry_trends(self, industry: str) -> Dict[str, Any]:
        """分析行业趋势"""
        trend_agent = self.agents["trend"]
        
        # 异步获取趋势分析
        response = await trend_agent.a_generate_reply(
            messages=[{
                "content": f"分析{industry}行业的最新趋势和技能需求",
                "role": "user"
            }]
        )
        
        return {
            "industry": industry,
            "analysis": response,
            "emerging_skills": self.extract_skills(response)
        }
    
    def extract_skills(self, text: str) -> List[str]:
        """从文本中提取技能关键词"""
        # 这里可以添加更复杂的NLP处理
        return list(set(text.split()))  # 简化处理
    
    def enhanced_skill_match(self, user_skills: List[str], job_skills: List[str]) -> float:
        """基于向量相似度的增强技能匹配"""
        # 获取技能向量
        user_embeddings = self.embedding_func(user_skills)
        job_embeddings = self.embedding_func(job_skills)
        
        # 计算相似度矩阵
        similarity_matrix = np.dot(user_embeddings, np.array(job_embeddings).T)
        
        # 取最大相似度作为匹配分数
        max_similarities = np.max(similarity_matrix, axis=1)
        return float(np.mean(max_similarities))
    
    def cluster_jobs(self, jobs_data: List[Dict]) -> Dict[int, List[Dict]]:
        """使用聚类分析岗位分组"""
        # 提取特征向量
        features = []
        for job in jobs_data:
            skills = job.get("skills", [])
            emb = np.mean(self.embedding_func(skills), axis=0)
            features.append(emb)
        
        # K-means聚类
        kmeans = KMeans(n_clusters=3, random_state=42)
        clusters = kmeans.fit_predict(features)
        
        # 分组结果
        clustered_jobs = {}
        for i, cluster_id in enumerate(clusters):
            if cluster_id not in clustered_jobs:
                clustered_jobs[cluster_id] = []
            clustered_jobs[cluster_id].append(jobs_data[i])
        
        return clustered_jobs

if __name__ == "__main__":
    # 初始化系统
    system = JobAgentSystem()
    
    # 示例使用
    async def demo():
        # 行业趋势分析
        trend = await system.analyze_industry_trends("IT")
        print(f"行业趋势分析: {trend}")
        
        # 技能匹配示例
        match_score = system.enhanced_skill_match(
            ["Python", "Machine Learning"],
            ["Python", "AI", "Data Science"]
        )
        print(f"技能匹配分数: {match_score:.2f}")
    
    asyncio.run(demo())